from torch.utils.data import Dataset
from constant import Constant
import numpy as np
import os


class TrafficDataset(Dataset):
    def __init__(self, train: bool = True):
        """初始化函数
        :param train: create train dataset if train, else create test dataset
        """
        name = 'train' if train else 'test'
        root = Constant.TRAIN_PATH if train else Constant.TEST_PATH

        data_path = os.path.join(root, f'{name}_data.npy')
        self.data = np.load(data_path)

        label_path = os.path.join(root, f'{name}_target.npy')
        self.label = np.load(label_path)

        tgt_in_path = os.path.join(root, f'{name}_tgt_in.npy')
        self.tgt_in = np.load(tgt_in_path)

    def __getitem__(self, item):
        return self.data[item], self.label[item], self.tgt_in[item]

    def __len__(self):
        return self.data.shape[0]
